from sklearn.feature_extraction.text import CountVectorizer as CV # 特征提取
from sklearn.model_selection import train_test_split
from sklearn.naive_bayes import MultinomialNB as NB # 多项式朴素贝叶斯分类器

# 加载SMS垃圾短信数据集
# as f将打开的文件对象赋值给变量f，以便后续对文件进行操作。
# sms遍历文件对象f中的每一行内容，对于每一行line，使用split('\t')方法按照制表符进行分割，将分割后的结果作为一个列表存储在sms中。
with open('./SMSSpamCollection.txt', 'r', encoding='utf8') as f:
    sms = [line.split('\t') for line in f]
# 使用zip(*sms)将列表中的标签和短信内容分离为两个列表y和x，分别表示短信的标签（是否为垃圾短信）和内容。
y, x = zip(*sms)  # x:短信内容，y:是或不是垃圾短信
# 数据预处理：将标签列表y中的标签转换为布尔值，表示是否为垃圾短信。具体地，将标签为'spam'的元素转换为True，其他标签转换为False。
y = [label == 'spam' for label in y]  # spam:垃圾短信
print(f"短信共有{len(y)}条")
print(f"垃圾短信(spam)有{y.count(True)}条")
print(f"正常短信(ham)有{y.count(False)}条")
# 数据集划分
x_train, x_test, y_train, y_test = train_test_split(x, y)
# 特征提取
counter = CV(token_pattern='[a-zA-Z]{2,}')
x_train = counter.fit_transform(x_train) # 文本向量化，将短信内容转换为特征向量
x_test = counter.transform(x_test) # 使用已有词汇表转换新文本
# print("x_train:", x_train)
# print(" x_test:", x_test)

# 朴素贝叶斯分类器
model = NB()
model.fit(x_train, y_train)
train_score = model.score(x_train, y_train)
test_score = model.score(x_test, y_test)
print("train score:", train_score)
print(" test score:", test_score)
